import chromadb
import dashscope
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
# chroma_client = chromadb.Client()
import os
def embed_query(content):
    resp = dashscope.TextEmbedding.call(
        model=dashscope.TextEmbedding.Models.text_embedding_v1,
        api_key="sk-8b48281f3ee3482fb107c4201d01da2b",  # 如果您没有配置环境变量，请将您的APIKEY填写在这里
        input=content)
    # print(resp)
    return resp

class aliEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        embeddings = []
        for i,text in enumerate(texts):
            print('生成Embedding::',i,text)
            embeddings.append(embed_query(text)['output']["embeddings"][0]['embedding'])
        # print(texts)
        return embeddings
from typing import Generator, List
import dashscope
from http import HTTPStatus


# 最多支持25条，每条最长支持2048tokens
DASHSCOPE_MAX_BATCH_SIZE = 25


def batched(inputs: List,
            batch_size: int = DASHSCOPE_MAX_BATCH_SIZE) -> Generator[List, None, None]:
    for i in range(0, len(inputs), batch_size):
        yield inputs[i:i + batch_size]
dashscope.api_key="sk-8b48281f3ee3482fb107c4201d01da2b"
def embed_with_list_of_str(inputs: List):
    result = None  # merge the results.
    batch_counter = 0
    embeddings = []
    for batch in batched(inputs):
        resp = dashscope.TextEmbedding.call(
            model=dashscope.TextEmbedding.Models.text_embedding_v1,
            input=batch)
        # print(resp)
        # if resp.status_code == HTTPStatus.OK:
        #     if result is None:
        #         result = resp
        #     else:
        
        for emb in resp['output']['embeddings']:
            emb['text_index'] += batch_counter
            # print(len(emb["embedding"]))
            embeddings.append(emb["embedding"])
        #     result.output['embeddings'].append(emb)
        # result.usage['total_tokens'] += resp.usage['total_tokens']
        # else:
        #     print(resp)
        batch_counter += len(batch)
    # print(result)
    return embeddings

class VectorDB():
    def __init__(self):
        self.client = chromadb.PersistentClient(path="./baoyandao")
        
    def getCollection(self,collectionName=""):
        self.collection = self.client.get_or_create_collection(name=collectionName, embedding_function=aliEmbeddingFunction(),
                metadata={"hnsw:space": "cosine"}) # l2 is the default)
    
    def getKDBList(self):
        return [{"title":c.name,"count":c.count()} for c in self.client.list_collections()]
    
    def delKDB(self,title):
        return self.client.delete_collection(name=title)
    
    def getKDBKList(self):#获取知识列表
        info = self.collection.get()
        print(info,type(info))
        info = list(zip(info["ids"],info["documents"],info["metadatas"]))
        info = [{"id":i[0],"k":i[1],"meta":i[2]} for i in info]
        return info
    
    def persistentEmbedding(self,data):
        embeddings = embed_with_list_of_str(data["documents"])
        # embeddings = [e["embedding"] for e in embeddings["embeddings"]]
        # print(embeddings,len(embeddings))
        dataList = list(zip(data["documents"],data["metadatas"],data["ids"],embeddings))
        print("开始持久化::>::",len(dataList))
        for idx,d in enumerate(dataList):
            print("写入数据库::",idx)
            self.collection.upsert(
                ids=[d[2]],
                embeddings=[d[3]],
                metadatas=[d[1]],
                documents=[d[0]],
            )
        # collection.add(
        #     documents=["taobao is a purchase web",
        #             "baidu is a search engine"],
        #     metadatas=[{"source": "doc1",'das':"dsa"},
        #             {"source": "doc2"}],
        #     ids=["id1",
        #         "id2"]
        # )
        # print(data["documents"])
        # print("client.heartbeat() ：：",self.client.heartbeat() )
        # print("开始持久化::",len(data["documents"]))
        # self.collection.add(
        #     documents=data["documents"],
        #     metadatas=data["metadatas"],
        #     ids=data["ids"]
        # )
        
        print("持久化完成::",len(data["documents"]))
        
        # print("--==--")

    def query(self,query,n_results=1):
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        # print(results)
        info = list(zip(results["ids"][0],results["documents"][0],results["metadatas"][0]))
        info = [{"id":i[0],"k":i[1],"meta":i[2]} for i in info]
        
        return info
    
import psycopg2
from psycopg2 import extras
class pgDB():
    def __init__(self):
# 建立数据库连接
        pass
        # self.con = psycopg2.connect(database="aiSearch",
        #                             user="postgres", 
        #                             password="xkkxkkxkk",
        #                             host="47.93.223.157",
        #                             port="5432")
        self.con=None

    def select(self,sql):
        self.con = psycopg2.connect(database="aiSearch",
                            user="postgres", 
                            password="xkkxkkxkk",
                            host="47.93.223.157",
                            port="5432")

        # 创建游标对象
        cur = self.con.cursor(cursor_factory=extras.DictCursor)

        # 使用游标执行SQL语句
        cur.execute(sql)

        # 获取查询结果
        data = cur.fetchall()
        data = [dict(row) for row in data]
        cur.close()
        self.con.close()
        # print(version)
        return data
        
    def modify(self,sql):
        # 创建游标对象
        self.con = psycopg2.connect(database="aiSearch",
                            user="postgres", 
                            password="xkkxkkxkk",
                            host="47.93.223.157",
                            port="5432")
        cur = self.con.cursor()

        # 使用游标执行SQL语句
        cur.execute(sql)

        # 获取查询结果
        self.con.commit()
        self.con.close()
        
    def persistentEmbedding(self,data):
        self.con = psycopg2.connect(database="aiSearch",
                            user="postgres", 
                            password="xkkxkkxkk",
                            host="47.93.223.157",
                            port="5432")

        # 创建游标对象
        cur = self.con.cursor(cursor_factory=extras.DictCursor)
        embeddings = embed_with_list_of_str(data["documents"])
        # embeddings = [e["embedding"] for e in embeddings["embeddings"]]
        # print(embeddings,len(embeddings))
        dataList = list(zip(data["documents"],data["fileName"],data["ids"],embeddings,data["oss_path"]))
        print("开始持久化::>::",len(dataList))
        # for idx,d in enumerate(dataList):
        #     print("写入数据库::",idx)
            # sql = '''
            # INSERT INTO knowledgeVALUES('{}','{}','{}','{}')
            # '''.format(dataList[idx][0],dataList[idx][3],dataList[idx][2],dataList[idx][1],)
            # self.modify(sql)
        # print(dataList[0])
        sql =  "INSERT INTO knowledge(content,file_name,kdb_id,vector,download_url)  " + "VALUES(%s, %s, %s,%s,%s)"
        
        cur = self.con.cursor()
        cur.executemany(sql, dataList)
        self.con.commit()
        self.con.close()
            # print(sql)
        print("持久化完成::",len(data["documents"]))
        
    def close(self):
        # 关闭连接
        # self.con.close()
        pass